#include "crypt.h"
#include <openssl/md5.h>
#include <openssl/rand.h>
#include <openssl/evp.h>
#include <string.h>
#include "misc.h"
#include "log.h"

#define ENC_INITED  1
#define DEC_INITED  2
#define ENC_STARTED 4
#define DEC_STARTED 8


void crypt_env_init() {
    OpenSSL_add_all_algorithms();
}

static int cipher_init(struct crypt *ctx, int enc) {
    EVP_CIPHER_CTX *evp;

    if (enc)
        evp = &ctx->enc;
    else
        evp = &ctx->dec;

    EVP_CIPHER_CTX_init(evp);
    if (!EVP_CipherInit_ex(evp, ctx->cipher, NULL, NULL, NULL, enc)) {
        LOGE("Cannot initialize cipher %s", EVP_CIPHER_name(ctx->cipher));
        return -1;
    }
    if (!EVP_CIPHER_CTX_set_key_length(evp, ctx->keylen)) {
        EVP_CIPHER_CTX_cleanup(evp);
        LOGE("Invalid key length: %"FMT_OF_SIZE, ctx->keylen);
        return -1;
    }

    EVP_CIPHER_CTX_set_padding(evp, 1);

    if (enc)
        ctx->flags |= ENC_INITED;
    else
        ctx->flags |= DEC_INITED;

    return 0;
}

int crypt_init(struct crypt *ctx, const char *ciphername, const char *passwd) {
    unsigned char iv[EVP_MAX_IV_LENGTH];

    ctx->flags = 0;
    ctx->decivlen = 0;

    if (!(ctx->cipher = EVP_get_cipherbyname(ciphername))) {
        LOGE("Cipher %s currently is not supported by PolarSSL library",
                ciphername);
        return -1;
    }

    ctx->keylen = EVP_BytesToKey(ctx->cipher, EVP_md5(), NULL,
            (const unsigned char *) passwd, strlen(passwd), 1, ctx->key, iv);
    if (ctx->keylen <= 0) {
        LOGE("Cannot generate key and IV");
        return -1;
    }

    if (cipher_init(ctx, 1))
        return -1;

    if (cipher_init(ctx, 0))
        return -1;

    return 0;
}

void crypt_destroy(struct crypt *ctx) {

    if (ctx->flags & ENC_INITED)
        EVP_CIPHER_CTX_cleanup(&ctx->enc);
    if (ctx->flags & DEC_INITED)
        EVP_CIPHER_CTX_cleanup(&ctx->dec);

    ctx->flags = 0;
    ctx->cipher = NULL;
}

static ssize_t encrypt_start(struct crypt *ctx, void *out, size_t outl) {
    if (!(ctx->flags & ENC_INITED)) {
        LOGE("crypt context not init.");
        return -1;
    }

    if (ctx->flags & ENC_STARTED)
        return 0;

    unsigned char iv[EVP_MAX_IV_LENGTH];
    size_t ivlen = EVP_CIPHER_CTX_iv_length(&ctx->enc);
    RAND_bytes(iv, ivlen);
    if (!EVP_CipherInit_ex(&ctx->enc, NULL, NULL, ctx->key, iv, 1)) {
        LOGE("Cannot set key and IV");
        return -1;
    }

    if (outl < ivlen)
        return -1;

    memcpy(out, iv, ivlen);

    ctx->flags |= ENC_STARTED;

    return ivlen;
}

ssize_t crypt_en(struct crypt *ctx, void *out, int outl, const void *in, int inl) {
    size_t total = 0;
    if (!(ctx->flags & ENC_STARTED)) {
        ssize_t n = encrypt_start(ctx, out, outl);
        if (n < 0)
            return -1;
        total += n;
        outl -= n;
        if (!outl)
            return total;
        out += n;
    }

    if (!EVP_CipherUpdate(&ctx->enc, out, &outl, in, inl))
        return -1;

    return total + outl;
}


static ssize_t decrypt_start(struct crypt *ctx, void *in, size_t inl) {
    if (!(ctx->flags & DEC_INITED)) {
        LOGE("crypt context not init.");
        return -1;
    }

    if (ctx->flags & DEC_STARTED)
        return 0;

    size_t ivlen = EVP_CIPHER_CTX_iv_length(&ctx->dec);

    size_t n = min(inl, ivlen - ctx->decivlen);
    memcpy(ctx->deciv + ctx->decivlen, in, n);
    ctx->decivlen += n;
    if (ctx->decivlen == ivlen) {
        if (!EVP_CipherInit_ex(&ctx->dec, NULL, NULL, ctx->key, ctx->deciv, 0)) {
            LOGE("Cannot set key and IV");
            return -1;
        }
        ctx->flags |= DEC_STARTED;
    }

    return n;
}

ssize_t crypt_de(struct crypt *ctx, void *out, int outl, void *in, int inl) {
    if (!(ctx->flags & DEC_STARTED)) {
        ssize_t n = decrypt_start(ctx, in, inl);
        if (n < 0)
            return -1;
        inl -= n;
        in += n;
        if (!inl)
            return 0;
    }

    if (!EVP_CipherUpdate(&ctx->dec, out, &outl, in, inl))
        return -1;

    return outl;
}
